﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Net.Sockets;
using System.Net;
using System.Threading;
using System.IO;

namespace LostLib.Net
{
	public class NetworkTransferException : IOException
	{
		public NetworkTransferException(string message, bool serverToClient)
			: base(message)
		{
			ServerToClient = serverToClient;
		}

		public bool ServerToClient { get; private set; }
	}

	[Serializable]
	public enum ExceptionReaction
	{
		Default,
		Abort,
		Ignore,
	}

	public class ExceptionHandlerEventArgs : EventArgs
	{
		public ExceptionHandlerEventArgs(Exception exception)
			: base()
		{
			if (exception == null) throw new ArgumentNullException("exception");

			Exception = exception;
			Reaction = ExceptionReaction.Default;
		}
		public ExceptionReaction Reaction { get; set; }
		public Exception Exception { get; private set; }
	}

	[Serializable]
	public enum SocksCommand : byte
	{
		Connect = 0x01,
		Bind = 0x02,
		UdpAssociate = 0x03,
	}

	[Serializable]
	public enum SocksAddressType : byte
	{
		IPv4 = 0x01,
		Domain = 0x03,
		IPv6 = 0x04,
	}

	[Serializable]
	public enum Socks4ServerReply : byte
	{
		OK = 0x5A,
		Fail = 0x5B,
		IdentdFailed = 0x5C,
		LoginFailure = 0x5D,
	}

	[Serializable]
	public enum Socks5ServerReply : byte
	{
		OK = 0x00,
		GeneralFailure = 0x01,
		DeniedByRules = 0x02,
		NetworkUnreachable = 0x03,
		HostUnreachable = 0x04,
		ConnectionRefused = 0x05,
		TtlExpired = 0x06,
		CommandNotSupported = 0x07,
		AddressNotSupported = 0x08,
	}

	public class Socks4ServerClient : SocksServerClient
	{
		internal Socks4ServerClient(TcpClient client)
			: base(client)
		{
		}

		const byte Socks4 = 4;
		const byte Zero = 0;
		const int MaxUserName = 256;
		const int BufferSize = 64 * 1024;

		private bool InitSocks4()
		{
			SocksCommand command;
			ushort port;
			byte[] ipv4;
			try
			{
				command = (SocksCommand)_source.ReadByte();
				port = (ushort)IPAddress.NetworkToHostOrder(_source.ReadInt16());
				ipv4 = BitConverter.GetBytes(_source.ReadInt32());
				byte[] username_buffer = new byte[MaxUserName];
				int i = 0;
				for (; i < MaxUserName; i++)
				{
					int chr = _source.Read();
					if (chr == -1)
						throw new IOException("Unexcepted end of stream. End of user name was excepted.");
					if (chr == 0)
						break;

					username_buffer[i] = (byte)chr;
				}
				if ((i == MaxUserName) && (_source.Read() != 0))
					throw new ApplicationException("User name too long.");
			} catch(Exception error)
			{
				OnTransferError(new ExceptionHandlerEventArgs(error));
				return false;
			}


			try
			{
				_remote = new TcpClient();
				_remote.Connect(new IPAddress(ipv4), port);
			} catch (SocketException error)
			{
				try
				{
					_dest.Write(Zero);
					_dest.Write((byte)Socks4ServerReply.Fail);
					_dest.Write(new byte[6]);
				} catch { }
				OnTransferError(new ExceptionHandlerEventArgs(error));
				return false;
			}
			try
			{
				_dest.Write(Zero);
				_dest.Write((byte)Socks4ServerReply.OK);
				_dest.Write(new byte[6]);
			} catch(Exception error)
			{
				OnTransferError(new ExceptionHandlerEventArgs(error));
				return false;
			}
			return true;
		}

		TcpClient _remote;
		NetworkStream _stream;
		BinaryReader _source;
		BinaryWriter _dest;

		public override void Run()
		{
			base.Run();

			using (_stream = Client.GetStream())
			using (_dest = new BinaryWriter(_stream))
			using (_source = new BinaryReader(_stream))
			{
				if (!InitSocks4()) return;

				ThreadPool.QueueUserWorkItem(new WaitCallback(ServerToClient));

				try
				{
					Client.NetworkTransfer(_stream, _remote.GetStream(), BufferSize);
					Client.Close();
					_remote.Close();
				} catch (Exception error)
				{
					var exception = new NetworkTransferException(error.Message, false);
					var args = new ExceptionHandlerEventArgs(exception);
					OnTransferError(args);
					Client.Close();
					_remote.Close();
					switch (args.Reaction)
					{
					case ExceptionReaction.Ignore:
					case ExceptionReaction.Abort:
						return;
					default:
						throw error;
					}
				}
			}
		}

		private void ServerToClient(object arg)
		{
			try
			{
				using (var remoteStream = _remote.GetStream())
				{

					_remote.NetworkTransfer(remoteStream, _stream, BufferSize);
					Client.Close();
					_remote.Close();

				}
			} catch (Exception error)
			{
				var exception = new NetworkTransferException(error.Message, false);
				var args = new ExceptionHandlerEventArgs(exception);
				OnTransferError(args);
				_remote.Close();
				Client.Close();
				switch (args.Reaction)
				{
				case ExceptionReaction.Ignore:
				case ExceptionReaction.Abort:
					return;
				default:
					throw error;
				}
			}
		}
	}

	public abstract class SocksServerClient : IDisposable
	{
		protected SocksServerClient(TcpClient client)
		{
			Client = client;
		}

		public TcpClient Client { get; protected set; }

		const byte Zero = 0;

		public event EventHandler<ExceptionHandlerEventArgs> TransferError;
		protected virtual void OnTransferError(ExceptionHandlerEventArgs args)
		{
			if (TransferError != null)
				TransferError(this, args);
		}

		public static SocksServerClient Start(TcpClient client)
		{
			var stream = client.GetStream();
			var source = new BinaryReader(stream, Encoding.ASCII);

			var ver = source.ReadByte();
			switch (ver)
			{
			//case 5:
			//    InitSocks5(source, dest);
			//    break;
			case 4:
				return new Socks4ServerClient(client);
			default:
				throw new NotSupportedException(string.Format("SOCKS protocol v{0} is not supported", ver));
			}
		}

		int _running;
		public bool Running
		{
			get
			{
				return _running != 0;
			}
		}
		/// <summary>
		/// 
		/// </summary>
		/// <remarks>
		/// You should call this in order to prevent mutiple calls to Run() in your child class.
		/// </remarks>
		public virtual void Run()
		{
			if (Interlocked.Increment(ref _running) != 1)
			{
				Interlocked.Decrement(ref _running);
				throw new InvalidOperationException("Client handler is already running");
			}
		}

		#region SOCKS5
		public const byte Socks5 = 5;
		private void InitSocks5(BinaryReader source, BinaryWriter dest)
		{
			var num_methods = source.ReadByte();
			byte auth_method = 0xFF;
			for (int i = 0; i < num_methods; i++)
			{
				var method = source.ReadByte();
				if (method == 0)
				{
					auth_method = method;
					break;
				}
			}
			dest.Write(Socks5);
			dest.Write(auth_method);

			switch (auth_method)
			{
			case 0:
				// NO AUTHORIZATION
				break;
			default:
				throw new NotSupportedException("None of suggested authorization methods are supported");
			}

			var ver = source.ReadByte();
			if (ver != 5)
				throw new NotSupportedException(string.Format("SOCKS protocol v{0} is not supported", ver));
			var command = (SocksCommand)source.ReadByte();
			source.ReadByte(); // RSV
			var addrType = (SocksAddressType)source.ReadByte();
			IPEndPoint destination = null;
			switch (addrType)
			{
			case SocksAddressType.IPv4:
				destination = GetIPv4Endpoint(source);
				break;
			case SocksAddressType.Domain:
				destination = GetDomainEndpoint(source);
				break;
			case SocksAddressType.IPv6:
				destination = GetIPv6Endpoint(source);
				break;
			}

			if (destination == null)
			{
				dest.Write(Socks5);
				dest.Write((byte)Socks5ServerReply.AddressNotSupported);
				dest.Write(Zero);
				dest.Write(Zero);
				dest.Write(0);
				dest.Write((ushort)0);
				throw new NotSupportedException("Address type not supported");
			}

			dest.Write(Socks5);
			dest.Write((byte)Socks5ServerReply.OK);
			dest.Write(Zero);

			throw new NotImplementedException();
		}
		#endregion SOCKS5

		#region Endpoints
		protected static IPEndPoint GetIPv4Endpoint(BinaryReader source)
		{
			byte[] ip = new byte[4];
			if (source.Read(ip, 0, 4) != 4)
				throw new IOException("Unexcepted end of stream. 4 bytes of IPv4 address was excepted");

			ushort port = source.ReadUInt16();
			return new IPEndPoint(new IPAddress(ip), port);
		}
		protected static IPEndPoint GetIPv6Endpoint(BinaryReader source)
		{
			byte[] ip = new byte[16];
			if (source.Read(ip, 0, 16) != 16)
				throw new IOException("Unexcepted end of stream. 16 bytes of IPv6 address was excepted");

			ushort port = source.ReadUInt16();
			return new IPEndPoint(new IPAddress(ip), port);
		}
		protected static IPEndPoint GetDomainEndpoint(BinaryReader source)
		{
			byte hostname_len = source.ReadByte();
			byte[] hostname_bytes = new byte[hostname_len];
			ushort port = source.ReadUInt16();
			if (source.Read(hostname_bytes, 0, hostname_len) != hostname_len)
				throw new IOException("Unexcepted end of stream. Host name was excepted");
			string hostname = Encoding.ASCII.GetString(hostname_bytes, 0, hostname_bytes.Length);
			var host_entry = Dns.GetHostEntry(hostname);
			return new IPEndPoint(host_entry.AddressList.First(), port);
		}
		#endregion

		//TcpClient _client;

		#region IDisposable Members

		public void Dispose()
		{
			Dispose(true);
			// This object will be cleaned up by the Dispose method.
			// Therefore, you should call GC.SupressFinalize to
			// take this object off the finalization queue
			// and prevent finalization code for this object
			// from executing a second time.
			GC.SuppressFinalize(this);
		}

		// Dispose(bool disposing) executes in two distinct scenarios.
		// If disposing equals true, the method has been called directly
		// or indirectly by a user's code. Managed and unmanaged resources
		// can be disposed.
		// If disposing equals false, the method has been called by the
		// runtime from inside the finalizer and you should not reference
		// other objects. Only unmanaged resources can be disposed.
		private void Dispose(bool disposing)
		{
			// Check to see if Dispose has already been called.
			if (Client != null)
			{
				// If disposing equals true, dispose all managed
				// and unmanaged resources.
				if (disposing)
				{
					// Dispose managed resources.
					Client.Close();
				}

				// Call the appropriate methods to clean up
				// unmanaged resources here.
				// If disposing is false,
				// only the following code is executed.

				Client = null;
			}
		}

		// Use C# destructor syntax for finalization code.
		// This destructor will run only if the Dispose method
		// does not get called.
		// It gives your base class the opportunity to finalize.
		// Do not provide destructors in types derived from this class.
		~SocksServerClient()
		{
			// Do not re-create Dispose clean-up code here.
			// Calling Dispose(false) is optimal in terms of
			// readability and maintainability.
			Dispose(false);
		}


		#endregion
	}
}
